#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Time    : 2025/4/25 16:35
@Author  : HZP
@File    : 3.RunnablePassthrough简化invoke调用.py
"""
from operator import itemgetter

from dotenv import load_dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI

load_dotenv()


def retrieval(query: str) -> str:
    """一个模拟的检索函数"""
    print("正在检索：", query)
    return "我是云信助手"


prompt = ChatPromptTemplate.from_template("""请根据用户的问题回答，可以参考对应的上下文进行生成。
<context>
{context}
</context>
用户的提问是: {query}""")
llm = ChatOpenAI(temperature=0.6, model="deepseek-r1:70b")

pars = StrOutputParser()

chain = RunnablePassthrough.assign(
    context=itemgetter("query"),
) | prompt | llm | pars
print(chain.invoke({"query": "你是谁？"}))
